{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "956ae8f6",
   "metadata": {},
   "source": [
    "# Relay ext 测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "edec8e0f",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "472a85cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "from tvm import relay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6c3f9b6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.register_func(\"relay.ext.test1\")\n",
    "def relay_ext_test(func):\n",
    "    return None\n",
    "\n",
    "mod = tvm.IRModule()\n",
    "shape = (relay.Any(), 25)\n",
    "dtype = \"float32\"\n",
    "\n",
    "# external function\n",
    "x = relay.var(\"x\", shape=shape, dtype=dtype)\n",
    "weight = relay.const(np.random.rand(5, 25).astype(\"float32\"), dtype=\"float32\")\n",
    "out = relay.nn.dense(x, weight)\n",
    "f1 = relay.Function([x], out)\n",
    "f1 = f1.with_attr(\"Primitive\", tvm.tir.IntImm(\"int32\", 1))\n",
    "f1 = f1.with_attr(\"Inline\", tvm.tir.IntImm(\"int32\", 1))\n",
    "f1 = f1.with_attr(\"Compiler\", \"test1\")\n",
    "f1 = f1.with_attr(\"global_symbol\", \"f1\")\n",
    "glb_f1 = relay.GlobalVar(\"f1\")\n",
    "mod[glb_f1] = f1\n",
    "mod = relay.transform.InferType()(mod)\n",
    "\n",
    "# Main function\n",
    "x = relay.var(\"x\", shape=shape, dtype=dtype)\n",
    "mod[\"main\"] = relay.Function([x], glb_f1(x))\n",
    "comp = relay.vm.VMCompiler()\n",
    "opt_mod, _ = comp.optimize(mod, target=\"llvm\")\n",
    "assert \"shape_func\" in opt_mod.astext(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "92eebb1f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@f1</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(?, <span style=\"color: #008000\">25</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(?, <span style=\"color: #008000\">25</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, Primitive<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Inline<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test1&quot;</span>, global_symbol<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;f1&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(?, <span style=\"color: #008000\">5</span>), float32] {\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>dense(<span style=\"color: #A2F; font-weight: bold\">%</span>x, meta[relay<span style=\"color: #A2F; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">0</span>] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">25</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, units<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(?, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x1: Tensor[(?, <span style=\"color: #008000\">25</span>), float32]) {\n",
       "  <span style=\"color: #A2F\">@f1</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x1)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "428c1552",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cannot find config for target=llvm -keys=cpu -mtriple=x86_64-unknown-linux-gnu, workload=('dense_pack.x86', ('TENSOR', (any_dim, 16), 'float32'), ('TENSOR', (any_dim, 16), 'float32'), None, 'float32'). A fallback configuration is used, which may bring great performance regression.\n"
     ]
    }
   ],
   "source": [
    "from tvm.relay.backend.vm import VMCompiler\n",
    "from tvm.relay.dataflow_pattern import is_op, wildcard\n",
    "from tvm.relay.loops import while_loop\n",
    "from tvm.relay.prelude import Prelude\n",
    "from tvm.relay.scope_builder import ScopeBuilder\n",
    "@tvm.register_func(\"relay.ext.test2\")\n",
    "def relay_ext_test(func):\n",
    "    return None\n",
    "\n",
    "data_shape = (relay.Any(), 16)\n",
    "weight_shape = (relay.Any(), 16)\n",
    "\n",
    "dense = relay.nn.dense(\n",
    "    relay.var(\"data\", shape=data_shape), relay.var(\"weight\", shape=weight_shape)\n",
    ")\n",
    "mod = tvm.IRModule.from_expr(dense)\n",
    "\n",
    "patterns = [(\"test.dense\", is_op(\"nn.dense\")(wildcard(), wildcard()))]\n",
    "passes = tvm.transform.Sequential(\n",
    "    [\n",
    "        relay.transform.MergeComposite(patterns),\n",
    "        relay.transform.AnnotateTarget([\"test2\"]),\n",
    "        relay.transform.PartitionGraph(),\n",
    "    ]\n",
    ")\n",
    "\n",
    "mod = passes(mod)\n",
    "\n",
    "compiler = VMCompiler()\n",
    "compiler.lower(mod, \"llvm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8d96354e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>data: Tensor[(?, <span style=\"color: #008000\">16</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(?, <span style=\"color: #008000\">16</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(?, <span style=\"color: #008000\">16</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(?, <span style=\"color: #008000\">16</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(?, ?), float32] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(?, <span style=\"color: #008000\">16</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(?, <span style=\"color: #008000\">16</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(?, <span style=\"color: #008000\">16</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(?, <span style=\"color: #008000\">16</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.dense_&quot;</span>, Composite<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test.dense&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(?, ?), float32] {\n",
       "    nn<span style=\"color: #A2F; font-weight: bold\">.</span>dense(<span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #A2F; font-weight: bold\">%</span>FunctionVar_0_1, units<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(?, ?), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(?, <span style=\"color: #008000\">16</span>), float32], Tensor[(?, <span style=\"color: #008000\">16</span>), float32]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(?, ?), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>data, <span style=\"color: #A2F; font-weight: bold\">%</span>weight) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(?, ?), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1571b0c9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aix",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
